#include "server/websocket_server.h"
#include "server/websocket_session.h"
#include "common/utils.h"
#include <iostream>
#include <cstring>
#include <unistd.h>
#include <arpa/inet.h>
#include <sstream>
#include <fstream>
#include <sys/stat.h>

namespace websocket {

WebSocketServer::WebSocketServer(int port) 
    : port_(port), server_socket_(-1), running_(false) {
}

WebSocketServer::~WebSocketServer() {
    stop();
}

bool WebSocketServer::start() {
    // 创建服务器socket
    server_socket_ = socket(AF_INET, SOCK_STREAM, 0);
    if (server_socket_ < 0) {
        std::cerr << "创建socket失败" << std::endl;
        return false;
    }
    
    // 设置socket选项
    int opt = 1;
    if (setsockopt(server_socket_, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt)) < 0) {
        std::cerr << "设置socket选项失败" << std::endl;
        close(server_socket_);
        return false;
    }
    
    // 绑定地址和端口
    sockaddr_in server_addr{};
    server_addr.sin_family = AF_INET;
    server_addr.sin_addr.s_addr = INADDR_ANY;
    server_addr.sin_port = htons(port_);
    
    if (bind(server_socket_, reinterpret_cast<sockaddr*>(&server_addr), sizeof(server_addr)) < 0) {
        std::cerr << "绑定端口 " << port_ << " 失败" << std::endl;
        close(server_socket_);
        return false;
    }
    
    // 开始监听
    if (listen(server_socket_, 10) < 0) {
        std::cerr << "监听失败" << std::endl;
        close(server_socket_);
        return false;
    }
    
    running_ = true;
    accept_thread_ = std::thread(&WebSocketServer::acceptConnections, this);
    
    std::cout << "WebSocket服务器启动成功，监听端口: " << port_ << std::endl;
    return true;
}

void WebSocketServer::stop() {
    if (!running_) return;
    
    running_ = false;
    
    // 关闭服务器socket
    if (server_socket_ >= 0) {
        close(server_socket_);
        server_socket_ = -1;
    }
    
    // 等待接受线程结束
    if (accept_thread_.joinable()) {
        accept_thread_.join();
    }
    
    // 关闭所有客户端连接
    {
        std::lock_guard<std::mutex> lock(sessions_mutex_);
        for (auto& session : sessions_) {
            session.second->close();
        }
        sessions_.clear();
    }
    
    std::cout << "WebSocket服务器已停止" << std::endl;
}

void WebSocketServer::acceptConnections() {
    while (running_) {
        sockaddr_in client_addr{};
        socklen_t client_len = sizeof(client_addr);
        
        int client_socket = accept(server_socket_, 
                                 reinterpret_cast<sockaddr*>(&client_addr), 
                                 &client_len);
        
        if (client_socket < 0) {
            if (running_) {
                std::cerr << "接受连接失败" << std::endl;
            }
            continue;
        }
        
        // 在新线程中处理客户端连接
        std::thread client_thread(&WebSocketServer::handleClient, this, client_socket);
        client_thread.detach();
    }
}

static std::string guess_mime(const std::string& path) {
    auto dot = path.find_last_of('.');
    std::string ext = (dot == std::string::npos) ? "" : path.substr(dot + 1);
    if (ext == "html") return "text/html; charset=UTF-8";
    if (ext == "css")  return "text/css; charset=UTF-8";
    if (ext == "js")   return "application/javascript; charset=UTF-8";
    if (ext == "json") return "application/json; charset=UTF-8";
    if (ext == "png")  return "image/png";
    if (ext == "jpg" || ext == "jpeg") return "image/jpeg";
    if (ext == "gif")  return "image/gif";
    return "text/plain; charset=UTF-8";
}

static bool read_file_to_string(const std::string& path, std::string& out) {
    std::ifstream ifs(path, std::ios::in | std::ios::binary);
    if (!ifs) return false;
    std::ostringstream ss;
    ss << ifs.rdbuf();
    out = ss.str();
    return true;
}

void WebSocketServer::handleClient(int client_socket) {
    char buffer[4096];
    ssize_t bytes_read = read(client_socket, buffer, sizeof(buffer) - 1);
    
    if (bytes_read <= 0) {
        close(client_socket);
        return;
    }
    
    buffer[bytes_read] = '\0';
    std::string request(buffer);

    // 获取请求行与路径
    // 形如：GET /path HTTP/1.1
    std::string path = "/";
    {
        std::istringstream req_stream(request);
        std::string method, uri, version;
        req_stream >> method >> uri >> version;
        if (!uri.empty()) path = uri;
    }

    // WebSocket 升级
    if (request.find("Upgrade: websocket") != std::string::npos) {
        if (performWebSocketHandshake(client_socket, request)) {
            auto session = std::make_shared<WebSocketSession>(client_socket);
            {
                std::lock_guard<std::mutex> lock(sessions_mutex_);
                sessions_[client_socket] = session;
            }
            session->setMessageCallback([this, session](const std::string& message) {
                if (message_handler_) message_handler_(session, message);
            });
            session->setCloseCallback([this, session]() {
                if (disconnection_handler_) disconnection_handler_(session);
                std::lock_guard<std::mutex> lock(sessions_mutex_);
                sessions_.erase(session->getSocket());
            });
            session->start();
            if (connection_handler_) connection_handler_(session);
            return;
        }
        // 握手失败：返回 400
        std::string bad = "HTTP/1.1 400 Bad Request\r\nConnection: close\r\nContent-Length: 0\r\n\r\n";
        sendData(client_socket, bad.c_str(), bad.size());
        close(client_socket);
        return;
    }

    // 普通 HTTP：提供简易静态资源（从 ./web 目录）
    // 路由规则：/ -> /login.html, /chat -> /index.html
    std::string mapped;
    if (path == "/" || path == "/index" || path == "/index.html") {
        mapped = "web/login.html";
    } else if (path == "/chat" || path == "/chat/") {
        mapped = "web/index.html";
    } else {
        // 去掉前导斜杠
        if (!path.empty() && path[0] == '/') mapped = "web" + path;
        else mapped = "web/" + path;
    }

    // 防止目录遍历
    if (mapped.find("..") != std::string::npos) {
        std::string resp = "HTTP/1.1 403 Forbidden\r\nConnection: close\r\nContent-Length: 0\r\n\r\n";
        sendData(client_socket, resp.c_str(), resp.size());
        close(client_socket);
        return;
    }

    std::string content;
    if (read_file_to_string(mapped, content)) {
        std::string mime = guess_mime(mapped);
        std::ostringstream hdr;
        hdr << "HTTP/1.1 200 OK\r\n"
            << "Content-Type: " << mime << "\r\n"
            << "Content-Length: " << content.size() << "\r\n"
            << "Connection: close\r\n\r\n";
        std::string header = hdr.str();
        sendData(client_socket, header.c_str(), header.size());
        sendData(client_socket, content.c_str(), content.size());
    } else {
        std::string body = "404 Not Found: " + path + "\n";
        std::ostringstream hdr;
        hdr << "HTTP/1.1 404 Not Found\r\n"
            << "Content-Type: text/plain; charset=UTF-8\r\n"
            << "Content-Length: " << body.size() << "\r\n"
            << "Connection: close\r\n\r\n";
        std::string header = hdr.str();
        sendData(client_socket, header.c_str(), header.size());
        sendData(client_socket, body.c_str(), body.size());
    }

    close(client_socket);
}

bool WebSocketServer::performWebSocketHandshake(int client_socket, const std::string& request) {
    // 提取Sec-WebSocket-Key
    size_t key_start = request.find("Sec-WebSocket-Key: ");
    if (key_start == std::string::npos) return false;
    
    key_start += 19; // "Sec-WebSocket-Key: "的长度
    size_t key_end = request.find("\r\n", key_start);
    if (key_end == std::string::npos) return false;
    
    std::string client_key = request.substr(key_start, key_end - key_start);
    std::string accept_key = generateWebSocketAcceptKey(client_key);
    
    // 构建握手响应
    std::string response = 
        "HTTP/1.1 101 Switching Protocols\r\n"
        "Upgrade: websocket\r\n"
        "Connection: Upgrade\r\n"
        "Sec-WebSocket-Accept: " + accept_key + "\r\n\r\n";
    
    return sendData(client_socket, response.c_str(), response.length());
}

std::string WebSocketServer::generateWebSocketAcceptKey(const std::string& client_key) {
    const std::string magic_string = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
    std::string combined = client_key + magic_string;
    
    unsigned char hash[20];
    utils::sha1(combined, hash);
    
    return utils::base64_encode(hash, 20);
}

void WebSocketServer::broadcast(const std::string& message) {
    std::lock_guard<std::mutex> lock(sessions_mutex_);
    for (auto& session : sessions_) {
        session.second->sendText(message);
    }
}

void WebSocketServer::broadcastTo(const std::vector<int>& sockets, const std::string& message) {
    std::lock_guard<std::mutex> lock(sessions_mutex_);
    for (int sock : sockets) {
        auto it = sessions_.find(sock);
        if (it != sessions_.end()) {
            it->second->sendText(message);
        }
    }
}

size_t WebSocketServer::getConnectionCount() const {
    std::lock_guard<std::mutex> lock(sessions_mutex_);
    return sessions_.size();
}

// 辅助函数：发送数据
bool WebSocketServer::sendData(int socket, const char* data, size_t length) {
    size_t sent = 0;
    while (sent < length) {
        ssize_t result = send(socket, data + sent, length - sent, 0);
        if (result <= 0) return false;
        sent += result;
    }
    return true;
}

} // namespace websocket